module hazard_Control(
    input clk,
    input rst_n,
    
    input [2:0] wd_sel,
    input rD1_en_ID,
    input rD2_en_ID,
    input rf_we_EX,
    input rf_we_MEM,
    input rf_we_WB,
    input [4:0] rR1_ID,
    input [4:0] rR2_ID,
    input [4:0] wR_EX,
    input [4:0] wR_MEM,
    input [4:0] wR_WB,
    input [31:0] wD_EX,
    input [31:0] wD_MEM,
    input [31:0] wD_WB,
    input npc_op,
    
    output forward_rD1,
    output forward_rD2,
    output reg stall_PC,
    output reg stall_IF_ID,
    output reg flush_IF_ID,
    output reg flush_ID_EX,

    output reg [31:0] rD1_f,
    output reg [31:0] rD2_f
    );
    
    // rD1 or rD2 is used in both ID and EX 
    wire rD1_ID_EX_hazard = (wR_EX == rR1_ID) & rf_we_EX & rD1_en_ID & (wR_EX != 5'b0);
    wire rD2_ID_EX_hazard = (wR_EX == rR2_ID) & rf_we_EX & rD2_en_ID & (wR_EX != 5'b0);
    
    // rD1 or rD2 is used in both ID and MEM 
    wire rD1_ID_MEM_hazard = (wR_MEM == rR1_ID) & rf_we_MEM & rD1_en_ID & (wR_MEM != 5'b0);
    wire rD2_ID_MEM_hazard = (wR_MEM == rR2_ID) & rf_we_MEM & rD2_en_ID & (wR_MEM != 5'b0);
    
    // rD1 is used in both ID and WB 
    wire rD1_ID_WB_hazard = (wR_WB == rR1_ID) & rf_we_WB & rD1_en_ID & (wR_WB != 5'b0);
    wire rD2_ID_WB_hazard = (wR_WB == rR2_ID) & rf_we_WB & rD2_en_ID & (wR_WB != 5'b0);

    // use forward or not
    assign forward_rD1 = rD1_ID_EX_hazard | rD1_ID_MEM_hazard | rD1_ID_WB_hazard;
    assign forward_rD2 = rD2_ID_EX_hazard | rD2_ID_MEM_hazard | rD2_ID_WB_hazard;

    // rD1 forwarding
    always @ (*) begin
        if(rD1_ID_EX_hazard)         rD1_f = wD_EX;
        else if(rD1_ID_MEM_hazard)   rD1_f = wD_MEM;
        else if(rD1_ID_WB_hazard)    rD1_f = wD_WB;
        else                         rD1_f = 32'b0;
    end

    // rD2 forwarding
    always @ (*) begin
        if(rD2_ID_EX_hazard)         rD2_f = wD_EX;
        else if(rD2_ID_MEM_hazard)   rD2_f = wD_MEM;
        else if(rD2_ID_WB_hazard)    rD2_f = wD_WB;
        else                         rD2_f = 32'b0;
    end
    
    // The last inst is load and it lead to data hazard
    wire load_hazard = (rD1_ID_EX_hazard | rD2_ID_EX_hazard) & (wd_sel[1:0] == 2'b01);
    
    // jump happens
    wire control_hz = npc_op;

    always @ (*) begin
        if(load_hazard) stall_PC = 1'b1;
        else            stall_PC = 1'b0;
    end

    always @ (*) begin
        if(load_hazard) stall_IF_ID = 1'b1;
        else            stall_IF_ID = 1'b0;
    end

    
    always @ (*) begin
        if(control_hz)  flush_IF_ID = 1'b1;
        else            flush_IF_ID = 1'b0;
    end

    
    always @ (*) begin
        if(load_hazard)     flush_ID_EX = 1'b1;
        else if(control_hz) flush_ID_EX = 1'b1;
        else                flush_ID_EX = 1'b0;
    end

endmodule
